"""
按周期计算收益率，简单方法
"""
import pandas as pd
class ror:
    def __init__(self,timeperiod=1):
        self.timeperiod=timeperiod+1
        self.data=[]
        self.timekey = None
    def oncalc(self,d,timekey=None):
        rst=0
        if pd.isna(d):
            if self.data:
                d = self.data[-1]
            else:
                return rst
        if timekey   is not None:
            if timekey == self.timekey:
                self.data[-1] = d
            else:
                self.data.append(d)
                self.timekey = timekey
        else:
            self.data.append(d)
        if len(self.data)>self.timeperiod:
            self.data.pop(0)
        if len(self.data)==self.timeperiod:
            rst=(d-self.data[0])/self.data[0]
        return rst
class ror_t2:
    def __init__(self,timeperiod=1):
        self.timeperiod=timeperiod+1
        self.data=[]
        self.timekey = None
    def oncalc(self,d,s,timekey=None):
        rst=0
        if pd.isna(d):
            if self.data:
                d = self.data[-1]
            else:
                return rst
        if timekey   is not None:
            if timekey == self.timekey:
                self.data[-1] = d
            else:
                self.data.append(d)
                self.timekey = timekey
        else:
            self.data.append(d)
        if len(self.data)>self.timeperiod:
            self.data.pop(0)
        if len(self.data)==self.timeperiod:
            if not pd.isna(s):
                rst=s/self.data[0]
        return rst
class ror_abs:
    def __init__(self,timeperiod=1):
        self.timeperiod=timeperiod+1
        self.data=[]
        self.timekey = None
    def oncalc(self,d,timekey=None):
        rst=0
        if pd.isna(d):
            if self.data:
                d=self.data[-1]
            else:
                return rst
        if timekey   is not None:
            if timekey == self.timekey:
                self.data[-1] = d
            else:
                self.data.append(d)
                self.timekey = timekey
        else:
            self.data.append(d)
        if len(self.data)>self.timeperiod:
            self.data.pop(0)
        if len(self.data)==self.timeperiod:
            rst=d-self.data[0]
        return rst

if __name__ == '__main__':
    from jili.core import load, save
    from research.calcor.calcors_graph import graph_calcor
    k1m = load(r"D:\data\future_k1m_tq\TA005\TA005_20200102.pkl")
    roc_stat0=ror(timeperiod=5)
    cc = [{'calc_cmd': 'ta', 'cmd': 'EMA', 'out': ['ema30'], 'input': {'price': 'close'}, 'arg': {'timeperiod': 30},
           'batch': 30},
          {'calc_cmd': 'ta', 'cmd': 'EMA', 'out': ['ema10'], 'input': {'price': 'close'}, 'arg': {'timeperiod': 10},
           'batch': 10},
          {'calc_cmd': 'function', 'function': roc_stat0.oncalc, 'out': ['roc_stat0'], 'input': ["ema10","timekey"]}]
    c = graph_calcor(cc)
    k1ms = k1m[:32]
    for t in k1ms:
        t = c.onbar(t)
    print(t)
    t1 = k1m[32]
    t1["timekey"] = t["timekey"]
    t1 = c.onbar(t1)
    print(t1)
    print("end")
